#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@author:hengk
@contact: hengk@foxmail.com
@datetime:2019-11-01 15:23
"""


from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from torchvision.transforms import transforms
import os,cv2
import json
import torch
from PIL import Image
import numpy as np

from datasets.noiser import Noiser
from datasets.image import Image

class ChartsCollate(object):
    """
    对每个batchsize中的数据进行合并,返回tensor类型的值
    """
    def __call__(self, items):
        batchsize = len(items)
        trans2norm = transforms.ToTensor()

        #根据最大宽 和 最大高，将图片进行scale
        max_h = max([item[0].shape[0] for item in items])
        max_w = max([item[0].shape[1] for item in items])


        images = torch.zeros((batchsize,3,max_h,max_w))
        labels = torch.Tensor(batchsize,10).long()

        for index,item in enumerate(items):
            images[index] = trans2norm(cv2.resize(item[0],(max_w,max_h)))
            labels[index] = item[1]

        return images,labels

class ChartSets(Dataset):
    """
    加载图片数据集
    """
    def __init__(self, images_dir, labels_dir,longer_size,is_train=True):
        self.images_dir_ = images_dir
        self.labels_dir_ = labels_dir
        self.image_names_ = os.listdir(images_dir)
        self.longer_size_ = longer_size
        self.is_train_ = is_train
        self.noiser_ = Noiser()
        self.image_ = Image()

    def __len__(self):
        return len(self.image_names_)

    def __getitem__(self, index):

        img_name = self.image_names_[index]
        label_name = img_name.split('.')[0] + ".json"
        with open(os.path.join(self.labels_dir_, label_name), encoding="utf-8") as fp:
            dicts = json.load(fp)
        del_set = ("is_legend_title", "is_wedge_text", "is_ylabel", "is_xlabel", "is_data_label")
        for e in del_set:
            del dicts[e]
        label = np.zeros((10))
        for index1, key in enumerate(dicts.keys()):
            if (dicts[key]):
                label[index1] = 1

        img = cv2.imread(os.path.join(self.images_dir_, img_name))
        if img is None:
            index = index + 1
            return self[index]

        # if(self.is_train_):
        #     img = self.noiser_.apply(img)
        #     img = self.image_.apply(img)

        h,w,_ = img.shape
        ratio = self.longer_size_/max(h, w)
        img = cv2.resize(img,(int(w*ratio), int(h*ratio)))
        label = torch.from_numpy(label)
        return img, label



if __name__ ==  "__main__":
    set = ChartSets("./data/train/images","./data/train/labels",512,False)
    collote = ChartsCollate()
    train_loader = DataLoader(
        set,
        batch_size=3,
        collate_fn=collote,
        shuffle=True,
        num_workers=3,
        drop_last=True,
        pin_memory=True)
    for index,(imgs,labels) in enumerate(train_loader):
        print(index)